import 'package:drift/drift.dart';
import 'package:path/path.dart' show url;
import 'package:recase/recase.dart';
import 'package:sqlparser/sqlparser.dart' as sql;

import '../analysis/options.dart';
import '../analysis/results/results.dart';
import '../utils/string_escaper.dart';
import 'import_manager.dart';
import 'queries/sql_writer.dart';

/// Manages a tree structure which we use to generate code.
///
/// Each leaf in the tree is a [StringBuffer] that contains some code. A
/// [Scope] is a non-leaf node in the tree. Why are we doing this? Sometimes,
/// we're in the middle of generating the implementation of a method and we
/// realize we need to introduce another top-level class! When passing a single
/// [StringBuffer] to the generators that will get ugly to manage, but when
/// passing a [Scope] we will always be able to write code in a parent scope.
class Writer extends _NodeOrWriter {
  late final Scope _root;
  late final TextEmitter _header;
  late final TextEmitter _imports;

  final DriftOptions options;
  final GenerationOptions generationOptions;

  TextEmitter get header => _header;

  TextEmitter get imports => _imports;

  @override
  Writer get writer => this;

  Writer(this.options, {required this.generationOptions}) {
    _root = Scope(parent: null, writer: this);
    _header = leaf();
    _imports = leaf();

    final preamble = options.preamble;
    if (preamble != null && generationOptions.writePreamble) {
      header.writeln(preamble);
    }
  }

  /// Returns the code generated by this [Writer].
  String writeGenerated() => _leafNodes(_root).join();

  Iterable<StringBuffer> _leafNodes(Scope scope) sync* {
    for (final child in scope._children) {
      if (child is TextEmitter) {
        yield child.buffer;
      } else if (child is Scope) {
        yield* _leafNodes(child);
      }
    }
  }

  Scope child() => _root.child();

  TextEmitter leaf() => _root.leaf();
}

final Uri modularSupport = Uri.parse('package:drift/internal/modular.dart');

abstract class _NodeOrWriter {
  Writer get writer;

  void _writeTagged(StringBuffer buffer, TaggedDartLexeme lexeme) {
    buffer.write(lexeme.lexeme);
  }

  AnnotatedDartCode generatedElement(DriftElement element, String dartName) {
    if (writer.generationOptions.isModular) {
      return AnnotatedDartCode.build(
          (b) => b.addGeneratedElement(element, dartName));
    } else {
      return AnnotatedDartCode([DartLexeme(dartName)]);
    }
  }

  AnnotatedDartCode modularAccessor(Uri driftFile) {
    final id = DriftElementId(driftFile, '(file)');

    return AnnotatedDartCode([
      DartTopLevelSymbol(
        ReCase(stripLeadingNumerics(url.basename(driftFile.path))).pascalCase,
        id.modularImportUri,
      ),
    ]);
  }

  AnnotatedDartCode companionType(DriftTable table) {
    if (table.nameOfCompanionClass case final customName?) {
      return generatedElement(table, customName);
    }

    final baseName = writer.options.useDataClassNameForCompanions
        ? table.nameOfRowClass
        : table.baseDartName;
    return generatedElement(table, '${baseName}Companion');
  }

  AnnotatedDartCode entityInfoType(DriftElementWithResultSet element) {
    return generatedElement(element, element.entityInfoName);
  }

  AnnotatedDartCode rowType(DriftElementWithResultSet element) {
    return AnnotatedDartCode.build((b) => b.addElementRowType(element));
  }

  AnnotatedDartCode rowClass(DriftElementWithResultSet element) {
    final existing = element.existingRowClass;
    if (existing != null) {
      return existing.targetClass ??
          (throw StateError('$element does not have a row class'));
    } else {
      return generatedElement(element, element.nameOfRowClass);
    }
  }

  /// Generates code that looks up [element] from an expression [database]
  /// evaluating to the attached database instance.
  ///
  /// This calls `resultSet()` with modular code and uses a direct field
  /// otherwise.
  AnnotatedDartCode referenceElement(
    DriftElementWithResultSet element,
    String database,
  ) {
    if (writer.generationOptions.isModular) {
      final infoType = entityInfoType(element);

      return AnnotatedDartCode.build((b) => b
        ..addSymbol('ReadDatabaseContainer', modularSupport)
        ..addText('($database).resultSet<')
        ..addCode(infoType)
        ..addText('>(${asDartLiteral(element.schemaName)})'));
    } else {
      return AnnotatedDartCode.text('$database.${element.dbGetterName}');
    }
  }

  /// Returns a Dart expression evaluating to the [converter].
  AnnotatedDartCode readConverter(AppliedTypeConverter converter,
      {bool forNullable = false}) {
    return AnnotatedDartCode.build((b) {
      final owningColumn = converter.owningColumn;
      final needsImplicitNullableVersion =
          forNullable && converter.canBeSkippedForNulls;
      final hasNullableVariantInField = owningColumn != null &&
          converter.canBeSkippedForNulls &&
          owningColumn.nullable;

      void addRegularConverter() {
        if (owningColumn != null) {
          b
            ..addCode(entityInfoType(owningColumn.owner))
            ..addText('.${converter.fieldName}');
        } else {
          // There's no field storing this converter, so evaluate it every time
          // it is used.
          b.addCode(converter.expression);
        }
      }

      switch ((needsImplicitNullableVersion, hasNullableVariantInField)) {
        case (false, _):
          addRegularConverter();
        case (true, false):
          b.addSymbol('NullAwareTypeConverter.wrap(', AnnotatedDartCode.drift);
          b.addCode(converter.expression);
          b.addText(')');
        case (true, true):
          b
            ..addCode(entityInfoType(converter.owningColumn!.owner))
            ..addText('.${converter.nullableFieldName}');
      }
    });
  }

  /// A suitable typename to store an instance of the type converter used here.
  AnnotatedDartCode converterType(AppliedTypeConverter converter,
      {bool makeNullable = false}) {
    // Write something like `TypeConverter<MyFancyObject, String>`
    return AnnotatedDartCode.build((b) {
      AnnotatedDartCode sqlDartType;

      switch (converter.sqlType) {
        case ColumnDriftType():
          sqlDartType =
              AnnotatedDartCode([dartTypeNames[converter.sqlType.builtin]!]);
        case ColumnCustomType(:final custom):
          sqlDartType = AnnotatedDartCode.type(custom.dartType);
      }

      final className = converter.alsoAppliesToJsonConversion
          ? 'JsonTypeConverter2'
          : 'TypeConverter';

      b
        ..addSymbol(className, AnnotatedDartCode.drift)
        ..addText('<')
        ..addDartType(converter.dartType)
        ..questionMarkIfNullable(makeNullable)
        ..addText(',')
        ..addCode(sqlDartType)
        ..questionMarkIfNullable(makeNullable || converter.sqlTypeIsNullable);

      if (converter.alsoAppliesToJsonConversion) {
        b
          ..addText(',')
          ..addDartType(converter.jsonType!)
          ..questionMarkIfNullable(
              makeNullable && !converter.jsonTypeIsNullable);
      }

      b.addText('>');
    });
  }

  AnnotatedDartCode dartType(HasType hasType) {
    return AnnotatedDartCode.build((b) => b.addDriftType(hasType));
  }

  /// The Dart type that matches the type of this column, ignoring type
  /// converters.
  ///
  /// This is the same as [dartType] but without type converters.
  AnnotatedDartCode variableTypeCode(HasType type,
      {bool? nullable, bool ignoreArray = false}) {
    if (type.isArray && !ignoreArray) {
      final inner =
          innerColumnType(type.sqlType, nullable: nullable ?? type.nullable);
      return AnnotatedDartCode([
        DartTopLevelSymbol.list,
        const DartLexeme('<'),
        ...inner.elements,
        const DartLexeme('>'),
      ]);
    } else {
      return innerColumnType(type.sqlType, nullable: nullable ?? type.nullable);
    }
  }

  /// The raw Dart type for this column, taking its nullability only from the
  /// [nullable] parameter.
  ///
  /// This type does not respect type converters or arrays.
  AnnotatedDartCode innerColumnType(ColumnType type, {bool nullable = false}) {
    return AnnotatedDartCode.build((b) {
      switch (type) {
        case ColumnDriftType():
          b.addTopLevel(dartTypeNames[type.builtin]!);
        case ColumnCustomType(:final custom):
          b.addDartType(custom.dartType);
      }

      if (nullable) {
        b.addText('?');
      }
    });
  }

  AnnotatedDartCode wrapInVariable(HasType column, AnnotatedDartCode expression,
      {bool ignoreArray = false}) {
    return AnnotatedDartCode.build((b) {
      b
        ..addTopLevel(DartTopLevelSymbol.drift('Variable'))
        ..addText('<')
        ..addCode(
            variableTypeCode(column, nullable: false, ignoreArray: ignoreArray))
        ..addText('>(');

      final converter = column.typeConverter;
      if (converter != null) {
        // apply type converter before writing the variable
        b
          ..addCode(readConverter(converter, forNullable: column.nullable))
          ..addText('.toSql(')
          ..addCode(expression)
          ..addText(')');
      } else {
        b.addCode(expression);
      }

      switch (column.sqlType) {
        case ColumnDriftType():
          break;
        case ColumnCustomType(:final custom):
          // Also specify the custom type since it can't be inferred from the
          // value passed to the variable.
          b
            ..addText(', ')
            ..addCode(custom.expression);
      }

      b.addText(')');
    });
  }

  String refUri(Uri definition, String element) {
    final prefix =
        writer.generationOptions.imports.prefixFor(definition, element);

    if (prefix == null) {
      return element;
    } else {
      return '$prefix.$element';
    }
  }

  /// References a top-level symbol exposed by the core `package:drift/drift.dart`
  /// library.
  String drift(String element) {
    return refUri(AnnotatedDartCode.drift, element);
  }

  String dartCode(AnnotatedDartCode code) {
    final buffer = StringBuffer();

    for (final lexeme in code.elements) {
      switch (lexeme) {
        case DartLexeme(:final lexeme):
          buffer.write(lexeme);
        case final TaggedDartLexeme tagged:
          _writeTagged(buffer, tagged);
        case DartTopLevelSymbol(importUri: final uri, :final lexeme):
          if (uri != null) {
            buffer.write(refUri(uri, lexeme));
          } else {
            buffer.write(lexeme);
          }
      }
    }

    return buffer.toString();
  }

  String sqlCode(sql.AstNode node, SqlDialect dialect) {
    return SqlWriter(writer.options, dialect: dialect, escapeForDart: false)
        .writeSql(node);
  }

  /// Builds a Dart expression writing the [node] into a Dart string.
  ///
  /// If the code for [node] depends on the dialect, the code returned evaluates
  /// to a `Map<SqlDialect, String>`. Otherwise, the code is a direct string
  /// literal.
  ///
  /// The boolean component in the record describes whether the code will be
  /// dialect specific.
  (String, bool) sqlByDialect(sql.AstNode node) {
    final dialects = writer.options.supportedDialects;

    if (dialects case [SqlDialect.sqlite]) {
      // Even if we only have a single dialect enabled, we should generate a
      // dialect-specific map if that dialect is not sqlite3. The reason is that
      // APIs in drift that aren't dialect-specific all assume sqlite3.
      return (
        SqlWriter(writer.options, dialect: dialects.single)
            .writeNodeIntoStringLiteral(node),
        false
      );
    }

    final buffer = StringBuffer();
    _writeSqlByDialectMap(node, buffer);
    return (buffer.toString(), true);
  }

  void _writeSqlByDialectMap(sql.AstNode node, StringBuffer buffer) {
    buffer.write('{');

    for (final dialect in writer.options.supportedDialects) {
      buffer
        ..write(drift('SqlDialect'))
        ..write(".${dialect.name}: '");

      SqlWriter(writer.options, dialect: dialect, buffer: buffer)
          .writeSql(node);

      buffer.writeln("',");
    }

    buffer.write('}');
  }
}

abstract class _Node extends _NodeOrWriter {
  final Scope? parent;

  _Node(this.parent);
}

/// A single lexical scope that is a part of a [Writer].
///
/// The reason we use scopes to write generated code is that some implementation
/// methods might need to introduce additional classes when written. When we can
/// create a new text leaf of the root node, this can be done very easily. When
/// we just pass a single [StringBuffer] around, this is annoying to manage.
class Scope extends _Node {
  final List<_Node> _children = [];
  @override
  final Writer writer;

  /// An arbitrary counter.
  ///
  /// This can be used to generated methods which must have a unique name-
  int counter = 0;

  /// The set of names already used in this scope. Used by methods like
  /// [getNonConflictingName] to prevent name collisions.
  final Set<String> _usedNames = {};

  Scope({required Scope? parent, Writer? writer})
      : writer = writer ?? parent!.writer,
        super(parent);

  DriftOptions get options => writer.options;

  GenerationOptions get generationOptions => writer.generationOptions;

  Scope get root {
    var found = this;
    while (found.parent != null) {
      found = found.parent!;
    }
    return found;
  }

  Scope child() {
    final child = Scope(parent: this);
    _children.add(child);
    return child;
  }

  TextEmitter leaf(
      {void Function(TaggedDartLexeme, StringBuffer)? writeTaggedDartCode}) {
    final child = TextEmitter(this, writeTaggedDartCode: writeTaggedDartCode);
    _children.add(child);
    return child;
  }

  /// Reserve a collection of names in this scope. See [getNonConflictingName]
  /// for more information.
  void reserveNames(Iterable<String> names) {
    _usedNames.addAll(names);
  }

  /// Returns a variation of [name] that does not conflict with any names
  /// already in use in this [Scope].
  ///
  /// If [name] does not conflict with any existing names then it is returned
  /// unmodified. If a conflict is detected then [name] is repeatedly passed to
  /// [modify] until the result no longer conflicts. Each result returned from
  /// this method is recorded in an internal set, so subsequent calls with the
  /// same name will produce a different, non-conflicting result.
  String getNonConflictingName(String name, String Function(String) modify) {
    while (_usedNames.contains(name)) {
      name = modify(name);
    }
    _usedNames.add(name);
    return name;
  }
}

class TextEmitter extends _Node {
  final StringBuffer buffer = StringBuffer();
  final void Function(TaggedDartLexeme, StringBuffer)? writeTaggedDartCode;

  @override
  final Writer writer;

  TextEmitter(Scope super.parent, {this.writeTaggedDartCode})
      : writer = parent.writer;

  @override
  void _writeTagged(StringBuffer buffer, TaggedDartLexeme lexeme) {
    if (writeTaggedDartCode case final function?) {
      function(lexeme, buffer);
    } else {
      super._writeTagged(buffer, lexeme);
    }
  }

  void write(Object? object) => buffer.write(object);

  void writeln(Object? object) => buffer.writeln(object);

  void writeUriRef(Uri definition, String element) {
    return write(refUri(definition, element));
  }

  void writeDriftRef(String element) => write(drift(element));

  void writeDart(AnnotatedDartCode code) => write(dartCode(code));

  void writeSql(sql.AstNode node,
      {required SqlDialect dialect, bool escapeForDartString = true}) {
    SqlWriter(
      writer.options,
      dialect: dialect,
      escapeForDart: escapeForDartString,
      buffer: buffer,
    ).writeSql(node);
  }

  void writeSqlByDialectMap(sql.AstNode node) {
    _writeSqlByDialectMap(node, buffer);
  }

  void stringLiteral(String contents) {
    return buffer.write(asDartLiteral(contents));
  }
}

/// Options that are specific to code-generation.
class GenerationOptions {
  /// Whether we're generating code to verify schema migrations.
  ///
  /// When non-null, we're generating from a schema snapshot instead of from
  /// source.
  final int? forSchema;

  /// Whether data classes should be generated.
  final bool writeDataClasses;

  /// Whether companions should be generated.
  final bool writeCompanions;

  /// Whether multiple files are generated, instead of just generating one file
  /// for each database.
  final bool isModular;

  /// Avoid pulling in user-code like type converters or `clientDefault`s.
  ///
  /// This is used internally when generating a `SchemaIsolate` used to export
  /// DDL statements.
  final bool avoidUserCode;

  /// Whether the [DriftOptions.preamble] should be written.
  ///
  /// We disable this when drift is generating part files picked up by the
  /// combining builder from `source_gen`, because users should set their
  /// preamble there.
  final bool writePreamble;

  final ImportManager imports;

  const GenerationOptions({
    required this.imports,
    this.forSchema,
    this.writeDataClasses = true,
    this.writeCompanions = true,
    this.isModular = false,
    this.avoidUserCode = false,
    this.writePreamble = true,
  });

  /// Whether, instead of generating the full database code, we're only
  /// generating a subset needed for schema verification.
  bool get isGeneratingForSchema => forSchema != null;
}

extension WriterUtilsForOptions on DriftOptions {
  String get fieldModifier => generateMutableClasses ? '' : 'final';
}

/// Adds an `this.` prefix is the [dartGetterName] is in [locals].
String thisIfNeeded(String getter, Set<String> locals) {
  if (locals.contains(getter)) {
    return 'this.$getter';
  }

  return getter;
}

extension on AnnotatedDartCodeBuilder {
  void questionMarkIfNullable(bool nullable) {
    if (nullable) addText('?');
  }
}
